import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from lppls import lppls
from datetime import datetime as dt

from data_source import get_data_by_meta
from utils import load_yaml_file


def train_lppl(data_kline):
    time = [pd.Timestamp.toordinal(dt.strptime(t1, '%Y-%m-%d')) for t1 in data_kline['Date']]
    price = np.log(data_kline['Adj Close'].values)
    # create observations array (expected format for LPPLS observations)
    observations = np.array([time, price])
    # set the max number for searches to perform before giving-up, the literature suggests 25
    # MAX_SEARCHES = 25
    # # instantiate a new LPPLS model with the data
    lppls_model = lppls.LPPLS(observations=observations)
    # # fit the model to the data and get back the params
    # tc, m, w, a, b, c, c1, c2, O, D = lppls_model.fit(MAX_SEARCHES)
    # lppls_model.plot_fit()
    res = lppls_model.mp_compute_nested_fits(workers=8, window_size=120, smallest_window_size=30, outer_increment=1,
                                             inner_increment=5, max_searches=25, )
    return lppls_model.compute_indicators(res)


def plot_bubble(res_df):
    fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, sharex=True, figsize=(18, 10))
    ord = res_df['time'].astype('int32')
    ts = [pd.Timestamp.fromordinal(d) for d in ord]
    # plot pos bubbles
    ax1_0 = ax1.twinx()
    ax1.plot(ts, res_df['price'], color='black', linewidth=0.75)
    ax1_0.plot(ts, res_df['pos_conf'], label='bubble indicator (pos)', color='red', alpha=0.5)
    # plot neg bubbles
    ax2_0 = ax2.twinx()
    ax2.plot(ts, res_df['price'], color='black', linewidth=0.75)
    ax2_0.plot(ts, res_df['neg_conf'], label='bubble indicator (neg)', color='green', alpha=0.5)

    ax1.grid(which='major', axis='both', linestyle='--')
    ax2.grid(which='major', axis='both', linestyle='--')
    # set labels
    ax1.set_ylabel('ln(p)')
    ax2.set_ylabel('ln(p)')

    ax1_0.set_ylabel('bubble indicator (pos)')
    ax2_0.set_ylabel('bubble indicator (neg)')

    ax1_0.legend(loc=2)
    ax2_0.legend(loc=2)

    plt.xticks(rotation=45)
    return fig



meta = load_yaml_file('config.yaml')
df_macro = get_data_by_meta(meta['meta'])

data_se = df_macro['布伦特原油现货']
data_se = data_se.dropna()
data_index = [str(s) for s in data_se.index]
time_se = [pd.Timestamp.toordinal(dt.strptime(t, '%Y-%m-%d')) for t in data_index]
price = np.log(data_se.values)
observations = np.array([time_se, price])
lppls_model = lppls.LPPLS(observations=observations)
res = lppls_model.mp_compute_nested_fits(
    workers=6, 
    window_size=120, 
    smallest_window_size=30, 
    outer_increment=1,
    inner_increment=5, 
    max_searches=25, 
    )
res_df = lppls_model.compute_indicators(res)
fig = plot_bubble(res_df)